# -*- coding: utf-8 -*-
''' 线性层 '''

import torch
import mymnist

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(50176, 10)

    def forward(self, x):
        return self.fc1(x)

net = Net()

for step, (images, labels) in enumerate(mymnist.test_loader):
    images = torch.flatten(images)
    output = net(images)
    print(output.shape)